import pyperclip
import pyotp
import time
import sys
import os
import shutil
import json
from argparse import ArgumentParser,Namespace
import threading
import base64
import binascii
import win32gui
from typing import *

import PyQt5
from PyQt5 import *
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *

namelist:dict[str,str]={
	"WindowTitle":"Github双因素验证密钥生成器",
	"Set":"设置",
	"DefaultKey":"000000",
	"Prompt":"0秒后失效",
	"CopiedDisplay":"已复制到剪贴板",
	"DialogTitle":"设置双因素验证SetupKey",
	"GetTip":"请输入SetupKey",
	"Tip":"提示",
	"Error":"错误",
	"Warning":"警告",
	"NonsetKey":"SetupKey不会改动。",
	"Nonbase32Digit":"不是合法的Base32字符串",
}
VERSION:int=1.0


def isKeyIllegality(key:str) -> bool:
	if not key or key.strip() == "":
		return False
	try:
		base64.b32decode(key)
	except binascii.Error:
		return False
	except ValueError:
		return False
	return True

class Dialog(QDialog):
	secret:str=str()
	available:bool=False

	def __init__(self:Self,parent=None) -> Self:
		super().__init__()
		self.main=parent
		self.setObjectName("GithubSetupKeyDialog")
		self.setWindowTitle(namelist["DialogTitle"])
		self.resize(393,104)
		self.setFont(QFont("Microsft YaHei UI",13))
		self.ui()
		self.retranslate()

	def ui(self:Self) -> None:
		self.Buttons:QDialogButtonBox=QDialogButtonBox(self)
		self.Buttons.setObjectName("Buttons")
		self.Buttons.setGeometry(QRect(120,60,141,32))
		self.Buttons.setOrientation(Qt.Horizontal)
		self.Buttons.setStandardButtons(QDialogButtonBox.Cancel | QDialogButtonBox.Ok)
		self.Buttons.setFont(QFont("Microsoft YaHei UI",12))

		self.GetSetupKey:QLineEdit=QLineEdit(self)
		self.GetSetupKey.setObjectName("GetSetupKey")
		self.GetSetupKey.setGeometry(QRect(140,10,241,31))
		self.GetSetupKey.setFont(QFont("Microsoft YaHei UI",12))

		self.SetupKey:QLabel=QLabel(self)
		self.SetupKey.setObjectName("SetupKey")
		self.SetupKey.setGeometry(QRect(10,10,131,31))
		self.SetupKey.setFont(QFont("Microsoft YaHei UI",12))

		self.Buttons.accepted.connect(self.Accept)
		self.Buttons.rejected.connect(self.Reject)

	def retranslate(self:Self) -> None:
		self.GetSetupKey.setPlaceholderText("".join("X" for _ in range(16)))
		self.SetupKey.setText(namelist["GetTip"]+":")

	def Accept(self:Self) -> None:
		self.secret=self.GetSetupKey.text()
		if not isKeyIllegality(self.secret):
			QMessageBox.critical(self,namelist["Error"],namelist["Nonbase32Digit"])
		self.available=True
		self.accept()

	def Reject(self:Self) -> None:
		self.available=False
		QMessageBox.warning(self,namelist["Warning"],namelist["NonsetKey"])
		self.reject()

class Main(QWidget):
	secret:str=str()
	reset:bool=False
	closed:bool=False
	whilethread:threading.Thread=None
	unset:bool=True

	def __init__(self:Self) -> Self:
		super().__init__()
		self.setObjectName("Window")
		self.setFixedSize(254,121)
		self.setFont(QFont("Microsoft YaHei UI",15))
		self.ui()
		self.retranslate()
	
	def ui(self:Self) -> None:
		self.Set:QPushButton=QPushButton(self)
		self.Set.setObjectName("Set")
		self.Set.setGeometry(QRect(170,5,71,71))

		self.Display:QLabel=QLabel(self)
		self.Display.setObjectName("Display")
		self.Display.setGeometry(QRect(10,10,151,41))
		self.Display.setFont(QFont("Microsoft YaHei UI",25))
		self.Display.setTextInteractionFlags(Qt.TextInteractionFlag.TextSelectableByMouse)
		self.Display.setAlignment(Qt.AlignCenter)

		self.Prompt:QLabel=QLabel(self)
		self.Prompt.setObjectName("Prompt")
		self.Prompt.setGeometry(QRect(10,80,231,41))
		self.Prompt.setFont(QFont("Microsoft YaHei UI",15))
		self.Prompt.setAlignment(Qt.AlignCenter)

		self.CopiedDidplay=QLabel(self)
		self.CopiedDidplay.setObjectName("CopiedDidplay")
		self.CopiedDidplay.setGeometry(QRect(10,50,151,31))
		self.CopiedDidplay.setFont(QFont("Microsoft YaHei UI",10))
		self.CopiedDidplay.setAlignment(Qt.AlignCenter)

		self.lineh1:QFrame=QFrame(self)
		self.lineh1.setObjectName("lineh1")
		self.lineh1.setGeometry(QRect(10,5,151,10))
		self.lineh1.setFrameShape(QFrame.HLine)
		self.lineh1.setFrameShadow(QFrame.Sunken)
		self.lineh2:QFrame=QFrame(self)
		self.lineh2.setObjectName("lineh2")
		self.lineh2.setGeometry(QRect(10,46,151,10))
		self.lineh2.setFrameShape(QFrame.HLine)
		self.lineh2.setFrameShadow(QFrame.Sunken)
		self.lineh3:QFrame=QFrame(self)
		self.lineh3.setObjectName("lineh3")
		self.lineh3.setGeometry(QRect(0,80,254,3))
		self.lineh3.setFrameShape(QFrame.HLine)
		self.lineh3.setFrameShadow(QFrame.Sunken)
		self.linev1:QFrame=QFrame(self)
		self.linev1.setObjectName("linev1")
		self.linev1.setGeometry(QRect(10,10,3,41))
		self.linev1.setFrameShape(QFrame.VLine)
		self.linev1.setFrameShadow(QFrame.Sunken)
		self.linev2:QFrame=QFrame(self)
		self.linev2.setObjectName("linev2")
		self.linev2.setGeometry(QRect(160,10,3,41))
		self.linev2.setFrameShape(QFrame.VLine)
		self.linev2.setFrameShadow(QFrame.Sunken)

		def setfn(self:Main):
			self.reset=True
			dialog:Dialog=Dialog(self)
			dialog.exec_()
			if dialog.available:
				self.secret=dialog.secret
			self.reset=False

		self.Set.clicked.connect(lambda:setfn(self))

	def retranslate(self:Self) -> None:
		self.setWindowTitle(namelist["WindowTitle"])
		self.Set.setText(namelist["Set"])
		self.Display.setText(namelist["DefaultKey"])
		self.Prompt.setText(namelist["Prompt"])
		self.CopiedDidplay.setText("")

	def closeEvent(self:Self,event:QCloseEvent | None) -> None:
		self.closed=True
		event.accept() if event else None

	def Start(self:Self) -> None:
		if not os.path.exists("config.json") or os.path.getsize("config.json") == 0:
			with open("config.json","w",encoding="utf-8") as f:
				dialog:Dialog=Dialog(self)
				dialog.exec_()
				if dialog.available:
					self.secret=dialog.secret
					json.dump({"key":self.secret},f)
					self.unset=False
		else:
			with open("config.json","r",encoding="utf-8") as f:
				self.secret=json.load(f)["key"]
				self.unset=False
		self.whilethread=threading.Thread(target=self.WhileDisplay)
		self.whilethread.start()

	def WhileDisplay(self:Self) -> NoReturn:
		while True:
			if self.closed:
				break
			if self.unset:
				self.Display.setText(namelist["DefaultKey"])
				self.Prompt.setText(namelist["Prompt"])
				self.CopiedDidplay.setText("")
				continue
			if not self.reset:
				code:int=pyotp.TOTP(self.secret).now()
				self.Display.setText(str(code))
				pyperclip.copy(code)
				self.CopiedDidplay.setText(namelist["CopiedDisplay"])
			for i in range(30):
				if self.closed or self.reset:
					break
				self.Prompt.setText(f"{str(30-i)}秒后失效")
				time.sleep(1)

def UiMain() -> NoReturn:
	handle:int=win32gui.GetForegroundWindow()
	win32gui.ShowWindow(handle,0)
	app:QApplication=QApplication(sys.argv)
	if not os.path.exists("ui.qss"):
		QMessageBox.critical(None,"Error","ui.qss not found!")
		win32gui.ShowWindow(handle,1)
		sys.exit(1)
	with open("ui.qss","r",encoding="utf-8") as qss:
		app.setStyleSheet(qss.read())
	window:Main=Main()
	window.show()
	window.Start()
	code:int=app.exec()
	win32gui.ShowWindow(handle,1)
	sys.exit(code)
		
def ConsoleMain() -> int:
	parser:ArgumentParser=ArgumentParser(description=f"Github双因素验证密钥生成器——Console版 v{VERSION}")
	parser.add_argument("-c","--console",nargs="?",metavar="",help="打开命令行模式")
	parser.add_argument("-s","--set",nargs="+",metavar="",help="设置密钥")
	parser.add_argument("-t","--temp-set",nargs="+",metavar="",help="临时使用密钥")
	
	args:Namespace=parser.parse_args()

	key:str=None
	path:str="config.json"
	prefix:str=""

	if args.set:
		if not isKeyIllegality(args.set[0]):
			print("\033[31m错误：\033[33m不是合法的Base32字符串\033[0m")
			return 1
		key=args.set[0]
		with open("config.json","w",encoding="utf-8") as cof:
			json.dump({"key":key},cof)

	else:
		if os.path.isfile(path):
			with open(path,"r",encoding="utf-8") as cof:
				key=json.load(cof)["key"]
				if not isKeyIllegality(key):
					try:
						key=input("Setup Key: ")
					except KeyboardInterrupt:
						print(f"\033[2A\n\r\033[2K\033[2K\033[31m已关闭\033[0m")
				if not isKeyIllegality(key):
					print("\033[31m错误：\033[33m不是合法的Base32字符串\033[0m")
					return 1
		else:
			key=input("Setup Key: ")
			if not isKeyIllegality(key):
				print("\033[31m错误：\033[33m不是合法的Base32字符串\033[0m")
				return 1
			with open("config.json","w",encoding="utf-8") as cof:
				json.dump({"key":key},cof)

	if args.temp_set:
		if not isKeyIllegality(args.temp_set[0]):
			print("\033[31m错误：\033[33m不是合法的Base32字符串\033[0m")
			return 1
		key=args.temp_set[0]
		prefix="临时 "

	try:
		while True:
			code:str=pyotp.TOTP(key).now()
			pyperclip.copy(code)
			for i in range(30):
				print(f"\033[33m{prefix}Setup Key: {key}\033[0m")
				print(f"\033[31m{code}\033[30m")
				print("\033[33m已复制到剪贴板\033[0m")
				print(f"\033[33m{30 - i}秒后失效\033[0m")
				print("\033[5A\r")
				time.sleep(1)
	except KeyboardInterrupt:
		print(f"\n\r\033[2K\n\r\033[2K\n\r\033[2K\r\033[3A\r\033[2K\033[31m已关闭\033[0m")

def Build() -> NoReturn:
	print("\033[33mStart building Simple 2FA for Github.\033[0m")
	print("\033[33mInstall required packages.\033[0m")
	code:int=os.system("pip install pywin32 pyperclip pyotp nuitka -i -i https://pypi.tuna.tsinghua.edu.cn/simple")
	if code != 0:
		print("\033[33mInstall failed!\033[0m")
		sys.exit(1)
	else:
		print("\033[33mInstalled.\033[0m")

	shutil.rmtree("out") if os.path.isdir("out") else None
	code:int=os.system(f"nuitka --mingw64 --no-deployment-flag=self-execution --enable-plugin=pyqt5 --standalone --output-filename=\"Simple 2FA for Github.exe\" --output-dir=out --show-progress --show-memory --follow-imports --follow-stdlib \"{sys.argv[0]}\"")
	if code == 0:
		print("\033[33mBuilt Simple 2FA for Github.\033[0m")
		sys.exit(0)
	print("\033[33mBuild failed!\033[0m")
	sys.exit(1)

if __name__ == "__main__":
	if len(sys.argv) == 1:
		UiMain()
	elif sys.argv[1] == "--build":
		if sys.argv[0].split(".")[-1] == "py":
			Build()
	else:
		sys.exit(ConsoleMain())
	


